#include "kernel_operator.h"
#include <type_traits>
#include "math.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_X, typename TYPE_Y>
class KernelGlobalAvgPool {

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TBuf<TPosition::VECCALC>calcBuf, xBuf, yBuf;   //

    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_Y> yGm;

    uint32_t minLength;
    
    uint32_t reduceNum; //进行reduce的一组数据的个数
    uint32_t groupNum;  //处理reduce的组数
    uint32_t totalLength; //总数据个数
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;    
    
public:
    __aicore__ inline KernelGlobalAvgPool(){}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t totalLength, uint32_t reduceNum, uint32_t groupNum, 
                                uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength) {

        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->minLength = 32 / sizeof(TYPE_X);
        this->totalLength = totalLength;
        this->reduceNum = reduceNum;
        this->groupNum = groupNum;

        this->tileNum = tileNum;
        this->tileLength = tileLength;
        this->lasttileLength = lasttileLength;


        xGm.SetGlobalBuffer((__gm__ TYPE_X*)x, totalLength);  //由于只有一个核心，所以偏移为0
        yGm.SetGlobalBuffer((__gm__ TYPE_X*)y, groupNum);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));   //会向上对齐32Byte

        pipe.InitBuffer(calcBuf, this->tileLength * sizeof(float));
        
        pipe.InitBuffer(xBuf, this->tileLength * sizeof(float));

        pipe.InitBuffer(yBuf, this->minLength * sizeof(float));
    }

    __aicore__ inline void Process() {
        
        int32_t loopCount = this->tileNum;
        for(int32_t i = 0; i < this->groupNum; i++)
        {
            //sum初始化0
            float sum = 0.0;
            for (int32_t j = 0; j < loopCount-1; j++)
            {
                CopyIn(i, j, tileLength);
                Compute(i, j, tileLength, sum);
            }
            CopyIn(i, tileNum-1, lasttileLength);
            Compute(i, tileNum-1, reduceNum - (tileNum-1) * tileLength, sum);
            // yGm.SetValue(i, (TYPE_X)(sum /reduceNum));
            // printf("sum : %f\n", sum);
            auto cursor = yGm[i].address_;
            *cursor = static_cast<TYPE_X>(sum /reduceNum);
        }    
    }


private:
    __aicore__ inline void CopyIn(int32_t i,  int32_t progress, uint32_t length) {
        LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
        // 本次copy的偏移量
        uint32_t start_index = i * reduceNum + progress * tileLength;
        //数据拷贝
        DataCopy(xLocal, xGm[start_index], length);
        inQueueX.EnQue(xLocal);
    }


    __aicore__ inline void Compute(int32_t i, int32_t progress, uint32_t length, float& sum) {
        LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        
        LocalTensor<float> yLocal_f = yBuf.Get<float>();
        LocalTensor<float> xLocal_f = xBuf.Get<float>();
        LocalTensor<float> wLocal = calcBuf.Get<float>();
        
        if constexpr(std::is_same_v<TYPE_X, half>)
        {   
            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);
        }
        else
        {
            xLocal_f = xLocal;
        }
        ReduceSum(xLocal_f, xLocal_f, wLocal, length);
        Adds(yLocal_f, xLocal_f, (float)sum, minLength);
        // yLocal_f.SetValue(0, yLocal_f.GetValue(0) + xLocal_f.GetValue(0));
        sum = yLocal_f.GetValue(0);
        inQueueX.FreeTensor(xLocal);
    }

};


extern "C" __global__ __aicore__ void global_avg_pool(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl

    KernelGlobalAvgPool<DTYPE_X, DTYPE_Y> op;
    op.Init(x, y, tiling_data.totalLength, tiling_data.reduceNum, tiling_data.groupNum, 
            tiling_data.tileNum, tiling_data.tileLength, tiling_data.lasttileLength);
    if(tiling_data.typeKey == 0)
    op.Process();
    else
    op.Process_FP16();
}